import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from configs import DEVICE, EPOCHS, INPUT_SIZE, HIDDEN_SIZE1, HIDDEN_SIZE2, OUTPUT_SIZE, MODEL_NAME, LEARNING_RATE
from model import FCNet, CNNet
from dataloader import prepare_train_data

def train_model(model,device,train_loader,optimizer,epoch):
    model.train()
    for batch_index,(data,label) in enumerate(train_loader):
        if MODEL_NAME == 'FCN' :
            data,label = data.reshape(-1,28*28).to(DEVICE),label.to(DEVICE)
        elif MODEL_NAME == 'CNN':
            data, label = data.to(device), label.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output,label)
        loss.backward()
        optimizer.step()
        if batch_index % 3000 == 0:
            print('Epoch[{}/{}]],Loss: {:.4f}'.format(epoch,EPOCHS,loss))

if MODEL_NAME == 'FCN' :
    model = FCNet(INPUT_SIZE, HIDDEN_SIZE1, HIDDEN_SIZE2, OUTPUT_SIZE).to(DEVICE)
elif MODEL_NAME == 'CNN' :
    model = CNNet(INPUT_SIZE, OUTPUT_SIZE).to(DEVICE)

optimizer = optim.Adam(model.parameters(),lr = LEARNING_RATE)
train_loader = prepare_train_data()
for epoch in range(1,EPOCHS + 1):
    train_model(model,DEVICE,train_loader,optimizer,epoch)
    torch.save(model.state_dict(),'./ckpt/'+MODEL_NAME+str(epoch)+'.pth')
    